﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Management;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Win32.SafeHandles;

namespace DiskInspector
{
    public static class DiskManager
    {
        private const int FileFlagNoBuffering = 0x20000000;
        private const int FileFlagSequentialScan = 0x08000000;
        private const int FileFlagOverlapped = 0x40000000;

        public const int SectorSize = 512;

        public static ManagementObject SelectedDisk;

        public static Stream GetStream()
        {
            if (SelectedDisk == null)
            {
                throw new InvalidOperationException("没有选择磁盘");
            }

            return GetStream(SelectedDisk);
        }

        public static Stream GetStream(ManagementObject disk)
        {
            var property = disk.Properties["Name"];
            var value = property.Value.ToString();
            return GetStream(value);
        }

        public static Stream GetStream(string fileName)
        {
            var handle = CreateFile(fileName, FileAccess.Read, FileShare.Read, IntPtr.Zero, FileMode.Open, FileFlagNoBuffering, IntPtr.Zero);
            var safeFileHandle = new SafeFileHandle(handle, true);
            return new FileStream(safeFileHandle, FileAccess.Read);
        }

        public static byte[] ReadSector(Stream disk, long offset)
        {
            return ReadSector(disk, offset, 1);
        }

        public static byte[] ReadSector(Stream disk, long sectorOffset, int sectorCount)
        {
            var count = SectorSize * sectorCount;
            var sectorData = new byte[count];
            var offset = 0;

            disk.Position = sectorOffset;
            int numberOfBytesRead;

            while ((numberOfBytesRead = disk.Read(sectorData, offset, count)) < count)
            {
                offset += numberOfBytesRead;
                count -= numberOfBytesRead;
            }

            return sectorData;
        }

        public static async Task<byte[]> ReadSectorAsync(Stream disk, long sectorOffset, int sectorCount)
        {
            return await ReadSectorAsync(disk, sectorOffset, sectorCount, new CancellationToken());
        }

        public static async Task<byte[]> ReadSectorAsync(Stream disk, long sectorOffset, int sectorCount, CancellationToken cancellationToken)
        {
            var count = SectorSize * sectorCount;
            var sectorData = new byte[count];
            var offset = 0;

            disk.Position = sectorOffset;

            do
            {
                var numberOfBytesRead = await disk.ReadAsync(sectorData, offset, count, cancellationToken);
                offset += numberOfBytesRead;
                count -= numberOfBytesRead;
            } 
            while (count > 0);

            return sectorData;
        }

        [DllImport("kernel32", SetLastError = true)]
        static extern IntPtr CreateFile(
            string FileName,
            FileAccess DesiredAccess,
            FileShare ShareMode,
            IntPtr SecurityAttributes,
            FileMode CreationDisposition,
            int FlagsAndAttributes,
            IntPtr hTemplate);
    }
}
